#include <iostream>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/segmentation/sac_segmentation.h>
#include <pcl/segmentation/extract_clusters.h>
#include <pcl/search/kdtree.h>

/**
 * @brief Remove sparse points to make it easier to be handled.
 * @param cloud
 * @param cloud_filtered
 * @author 1160300719 殷浩然
 */
void Remove_Sparse_Points(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filtered){
    pcl::PCDWriter writer;
    pcl::StatisticalOutlierRemoval<pcl::PointXYZ> sor;
    sor.setInputCloud(cloud);
    sor.setMeanK(5);
    sor.setStddevMulThresh(0.5);
    sor.filter(*cloud_filtered);
    writer.write("../pcd/cloud_filter.pcd", *cloud_filtered, false);
}

/**
 * @brief Remove the plane and save the plane for the next step.
 * @param cloud_filtered
 * @param inlier
 * @param coefficients
 * @param cloud_p
 * @param cloud_np
 * @author 1160300719 殷浩然
 */
void Remove_Plane(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filtered, pcl::PointIndices::Ptr inlier, pcl::ModelCoefficients::Ptr coefficients, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_p, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_np){
    pcl::ExtractIndices<pcl::PointXYZ> extract;
    pcl::SACSegmentation<pcl::PointXYZ> seg;
    pcl::PCDWriter writer;
    seg.setOptimizeCoefficients(true);
    seg.setModelType(pcl::SACMODEL_PLANE);
    seg.setDistanceThreshold(0.01);
    seg.setInputCloud(cloud_filtered);
    seg.segment(*inlier, *coefficients);
    extract.setInputCloud(cloud_filtered);
    extract.setIndices(inlier);
    extract.setNegative(false);
    extract.filter(*cloud_p);
    extract.setNegative(true);
    extract.filter(*cloud_np);
    writer.write("../pcd/cloud_plane.pcd", *cloud_p, false);
    writer.write("../pcd/cloud_unplane.pcd", *cloud_np, false);
}

/**
 * @brief Use Euclidean Cluster to separate the rest part.
 * @param cloud_np
 * @param cluster_indices
 * @author 1160300719 殷浩然
 */
void Euclidean_Cluster(pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_np, std::vector<pcl::PointIndices> cluster_indices){
    pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>);
    pcl::ExtractIndices<pcl::PointXYZ> extract;
    pcl::PCDWriter writer;
    tree->setInputCloud(cloud_np);
    pcl::EuclideanClusterExtraction<pcl::PointXYZ> ec;
    ec.setClusterTolerance(0.004);
    ec.setMinClusterSize(100);
    ec.setMaxClusterSize(25000);
    ec.setSearchMethod(tree);
    ec.setInputCloud(cloud_np);
    ec.extract(cluster_indices);
    for(unsigned int i=0; i<cluster_indices.size(); i++){
        pcl::PointCloud<pcl::PointXYZ>::Ptr test_cloud(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::PointIndices::Ptr temp_pointer(new pcl::PointIndices);
        *temp_pointer = cluster_indices.at(i);
        extract.setInputCloud(cloud_np);
        extract.setIndices(temp_pointer);
        extract.setNegative(false);
        extract.filter(*test_cloud);
        char ch[20];
        sprintf(ch, "../pcd/rest_part_%d.pcd", i);
        writer.write(ch, *test_cloud, false);
    }
}

/**
 * @brief The main process of the programme.
 * @return The state of the programme.
 * @author 1160300719 殷浩然
 */
int main()
{
    //Params
    std::vector<pcl::PointIndices> cluster_indices;
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filtered(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_p(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_np(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::PointCloud<pcl::PointXYZ>::Ptr add_cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::ModelCoefficients::Ptr coefficients(new pcl::ModelCoefficients);
    pcl::PointIndices::Ptr inlier(new pcl::PointIndices);

    //Read source
    pcl::PCDReader reader;
    int state = reader.read("../pcd/out.pcd", *cloud);
    if(state == -1){
        PCL_ERROR("File is not open.");
        return -1;
    }

    Remove_Sparse_Points(cloud, cloud_filtered);
    Remove_Plane(cloud_filtered, inlier, coefficients, cloud_p, cloud_np);
    Euclidean_Cluster(cloud_np, cluster_indices);

    return 0;
}
